import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from FlagEmbedding import BGEM3FlagModel
from pymilvus import MilvusClient
import numpy as np

# 1. 初始化BGE-M3模型
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)

# 2. 初始化Milvus Lite客户端
client = MilvusClient("milvus_demo.db")

# 3. 准备数据
corpus = [
    "自然语言处理是人工智能的一个重要领域",
    "大型语言模型在许多NLP任务中表现出色",
    "计算机视觉处理图像和视频数据",
    "推荐系统基于用户行为预测偏好",
    "信息检索关注如何高效获取相关信息"
]

# 4. 创建Milvus集合
collection_name = "bge_m3_docs"
dim = 1024  # BGE-M3的向量维度

if client.has_collection(collection_name):
    client.drop_collection(collection_name)

# 创建集合
client.create_collection(
    collection_name=collection_name,
    dimension=dim,
    primary_field_name="id",
    id_type="int",
    primary_field="id",
    auto_id=True,
    vector_field_name="embedding",
    metric_type="IP"
)

# 5. 生成向量并插入Milvus
def get_embeddings(texts):
    return model.encode(texts, batch_size=128, max_length=8192)['dense_vecs']

# 插入文档
embeddings = get_embeddings(corpus)
for text, embedding in zip(corpus, embeddings):
    client.insert(
        collection_name=collection_name,
        data={"text": text, "embedding": embedding}
    )

# Milvus Lite自动创建索引，无需手动操作

# 7. 检索函数
def search(query, top_k=3):
    # 生成查询向量
    query_embedding = get_embeddings([query])[0]
    
    # 准备搜索参数
    search_params = {
        "metric_type": "IP",
        "params": {"nprobe": 10}
    }
    
    # 执行搜索
    results = client.search(
        collection_name=collection_name,
        data=[query_embedding],
        limit=top_k,
        output_fields=["text"]
    )
    
    # 格式化结果
    ret = []
    for hit in results[0]:
        ret.append({
            "text": hit["text"],
            "score": hit["distance"],
            "id": hit["id"]
        })
    return ret

# 8. 示例查询
query = "什么技术用于文本分析？"
print(f"查询: {query}")
print("搜索结果:")
for result in search(query):
    print(f"ID: {result['id']} | 相似度: {result['score']:.4f} | 文档: {result['text']}")

# 9. 清理
client.close()